import torch
from torch import nn
from data import dataset
import numpy as np


class Logistic(nn.Module):
    def __init__(self):
        super(Logistic, self).__init__()
        self.fc = nn.Linear(2, 1)
        self.sigmod = nn.Sigmoid()

    def forward(self, x):
        x = self.fc(x)
        x = self.sigmod(x)
        return x


model = Logistic()
# 优化器
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3, momentum=0.9)
# 损失函数
criterion = nn.BCELoss()

data, labels = tuple(map(torch.tensor,  dataset.loadDataSet("./data/data.txt")))
for i in range(10000):
    optimizer.zero_grad()
    out = model(data)
    loss = criterion(out, labels.unsqueeze(dim=1))
    loss.backward()
    optimizer.step()

weight = model.state_dict()['fc.weight'].numpy()
b = model.state_dict()['fc.bias'].numpy()
print(weight[0], b)
plt, ax = dataset.Drawdata()
x = np.arange(-4.0, 4.0, 0.1)
y = -(b + weight[0, 0] * x) / weight[0, 1]
ax.plot(x, y)
plt.show()

